import idaapi
import idautils
import idc
import shelve
import time
from DisTools.Config import *


# 不能使用yaml,需要在ida的内置python里面增加安装包才支持
extract_error_log = logs["errorLog"]
graphInfLog = logs["graphInfLog"]
dataSavePath = data["dataSavePath"]
arithmetric=["add","adc","inc","daa","sub","sbb","dec","cmp","das","mul","imul","div","idiv"]
dataTrans=["mov","push","pop","xchg","xlat","cwd","lea","lds","les"]
controlTrans=["call","ret","retn","jmp","loop","loope","loopne","int","into","iret","jz",\
    "je","jnz","jne","jo","jno","jc","jnc","js","jns","jb","jnae","jnb","jae","ja",\
        "jnbe","jna","jbe","jl","jnge","jnl","jge","jg","jnle","jng","jle"]
sample_extrace = data["extracAllData"]


def function_call_extract(func,callees,callers):
    """
        保存函数间的调用关系
    """
    func_name = idc.get_func_name(func)  
    print("{}被调用".format(func_name))

    # 同一函数可能会多次调用另一个函数
    for ref_ea in idautils.CodeRefsTo(func, 0):    
        caller_name = idc.get_func_name(ref_ea)
        print("    {}--{}".format(hex(ref_ea),idc.GetDisasm(ref_ea)))
        print("    {}调用".format(caller_name))
        # 移除未识别出函数的情况
        if caller_name!="":
            callees[caller_name].append(func_name) 
            callers[func_name].append(caller_name)

def function_extract(func,callees,callers):
    """
        #增加提取函数属性
        #增加保存函数汇编指令
    """
    func_name = idc.get_func_name(func)    
    temp=idaapi.get_func(func)
    start=temp.start_ea
    end=temp.end_ea
    cur_addr=start
    discode=[]
    emb=[0]*8
    #按照Gemini的方式提取，中心性感觉复杂度很高，目前没有实现
    while cur_addr<end:
        #保存汇编指令
        discode.append("{}--{}".format(hex(cur_addr),idc.GetDisasm(cur_addr)))
    #函数内特征
        #指令数量
        emb[4]+=1

        #字符串常量,提取存在问题,提取不出来
        insn = idaapi.insn_t()
        if idaapi.decode_insn(insn, cur_addr):  
            for operand in insn.ops:
                # 判断操作数是否为立即数操作数
                if operand.type == idaapi.o_imm:
                    # 获取立即数的值
                    immediate_value = operand.value
                    if idc.is_strlit(immediate_value):
                        emb[0]+=1
                        print("字符常量{}".format(idc.GetDisasm(cur_addr)))
                if operand.type == idaapi.o_mem:
                    mem_address = operand.addr
                    print("常量{}".format(idc.GetDisasm(cur_addr)))
                    if idc.is_strlit(mem_address):
                        emb[0]+=1
                        print("字符常量{}".format(idc.GetDisasm(cur_addr)))

        #数字常量  && 字符串常来给你
        insn = idaapi.insn_t()
        if idaapi.decode_insn(insn, cur_addr) > 0:  
            for op in insn.ops:
                if op.type == idaapi.o_imm:
                    emb[1]+=1   
                    # print("数值常量{}".format(idc.GetDisasm(cur_addr)))

        #转移指令
        if(idc.print_insn_mnem(cur_addr) in dataTrans):
            emb[2]+=1
        #控制指令
        elif(idc.print_insn_mnem(cur_addr) in controlTrans):
            emb[3]+=1
        #算术指令
        elif(idc.print_insn_mnem(cur_addr) in arithmetric):
            emb[5]+=1
        #下一条指令起点
        cur_addr=idc.next_head(cur_addr,end)
    #函数间特征
    #后代数量
    emb[6]=len(callers[func_name])
    #上一辈数量
    if func_name in callees:
        emb[7]=len(callees[func_name])
    return emb,discode

def cg_extract(funcs_id, callees, func_num):
    cg_adjmat=[]
    for i in range(func_num):
        cg_adjmat.append([0]*func_num)
    for key in callees.keys():
        for calling in callees[key]:
            #有边则加一
            # cg_adjmat[funcs_id[key]][funcs_id[calling]]+=1
            cg_adjmat[funcs_id[calling]][funcs_id[key]]+=1
    return cg_adjmat

def BB_extract(func):
    cnt = 0
    f = idaapi.FlowChart(idaapi.get_func(func))
    cfg_adjmat = []
    cfg_attr=[]
    cfg_dis=[]
    
    for block in f:
        cfg_row =[0]*f.size
        succNum=0
        predNum=0
        for succ_block in block.succs():
            cfg_row[succ_block.id] = 1
            succNum+=1
        for i in block.preds():
            predNum+=1

        cur_addr=block.start_ea
        emb=[0]*8
        emb[6]=succNum
        emb[7]=predNum
        temp=[]
        while cur_addr<block.end_ea: #遍历汇编代码
            temp.append("{}--{}".format(hex(cur_addr),idc.GetDisasm(cur_addr)))
            if idaapi.is_strlit(idaapi.get_flags(cur_addr)):
                emb[0] += 1
            emb[4]+=1
            if(idc.get_operand_type(cur_addr,0)==5 or idc.get_operand_type(cur_addr,1)==5 or idc.get_operand_type(cur_addr,2)==5): 
                emb[1]+=1
            if(idc.print_insn_mnem(cur_addr) in dataTrans):
                emb[2]+=1
            elif(idc.print_insn_mnem(cur_addr) in controlTrans):
                emb[3]+=1
            elif(idc.print_insn_mnem(cur_addr) in arithmetric):
                emb[5]+=1
            cur_addr=idc.next_head(cur_addr,block.end_ea)
        cfg_dis.append(temp)
        cfg_attr.append(emb)
        cfg_adjmat.append(cfg_row)
    return cfg_adjmat,cfg_attr,cfg_dis
                                
def get_imports_name():
    imports={}
    current=""
    def callback(ea,name,oridinal):
        imports[current].append((hex(ea),name,oridinal))
        return True
    muduleNum = idaapi.get_import_module_qty()
    for i in range(0, muduleNum):
        module_name = idaapi.get_import_module_name(i)
        print("模块名称:", module_name)
        current=module_name
        imports[module_name]=[]
        idaapi.enum_import_names(i, callback)
    return imports

def controller():
    funcs_id = dict() #函数--》id
    callees = dict()  #被调用者 
    callers = dict()  #被调用者 
    funcType=dict()
    func_num = 0
    func_id = 0

    funcs = idautils.Functions()#遍历函数的指针
    funcs1 = idautils.Functions()
    funcs_iterator = idautils.Functions()
    
    # 获得函数及其对应的下标（id）
    for i in funcs_iterator:
        func_name = idc.get_func_name(i)
        flag = "local"
        #判断是否为静态链接库代码,
        # ida中为蓝色标注为静态连接库函数，粉色的一般为调用动态导入函数的简单thunk函数
        flags=idc.get_func_flags(i)
        if(flags & FUNC_LIB):
            flag = "static link"
        funcType.update({func_id:flag})
        funcs_id.update({func_name:func_id})
        func_num += 1
        func_id += 1
        callers[func_name]=[]
        callees[func_name]=[]
    
    # 动态导入函数一般不在Functions中，要获取导入表中各个api的名称才行
    apiName=get_imports_name()
    flag = "dynamic import"
    # 加入动态导入函数
    for moduleName in apiName:
        for addr,tarname,temp in apiName[moduleName]:
            addr = idc.get_name_ea_simple(tarname)#二进制的api地址
            n_id=func_id
            # Function中同名的Thunk函数，修改函数类别即可
            if tarname in funcs_id:
                n_id=funcs_id[tarname]
                funcType.update({n_id:flag})
            #不存在api函数，添加！
            else:
                funcType.update({func_id:flag})
                funcs_id.update({tarname:func_id})
                func_num += 1
                func_id += 1
                callers[tarname]=[]
                callees[tarname]=[]

    #分别用来保存cfg图，对应的属性，cg的属性和每个函数的汇编代码
    cfg,attr,cattr,discodes,cdis=[],[],[],[],[]
    for i in range(len(funcs_id)):
        cfg.append([])
        attr.append([])
        cattr.append([])
        discodes.append([])
        cdis.append([])

    # 构建函数间的调用关系
    for f in funcs:
        func_name = idc.get_func_name(f)
        #移除动态导入函数的thunk函数，避免重复计算
        if funcType[funcs_id[func_name]]!="dynamic import":
            #callers 保存调用key的所有函数，callees 保存key调用的所有函数  
            function_call_extract(f, callees,callers)

    #api的函数调用也需要添加进来
    for moduleName in apiName:
        for addr,tarname,temp in apiName[moduleName]:
            addr = idc.get_name_ea_simple(tarname)
            print(tarname)
            for i in idautils.CodeRefsTo(addr,0):
                caller_name = idc.get_func_name(i)
                # 移除未识别出函数的情况
                if caller_name !="":
                    print("    {}--{}".format(hex(i),idc.GetDisasm(i)))
                    print("    {}调用".format(caller_name))
                    callees[caller_name].append(tarname) 
                    callers[tarname].append(caller_name)
    
    for f in funcs1:        
        func_name = idc.get_func_name(f)     
        # 提取函数特征向量
        # 只需要提取functions里面的函数，动态连接的函数无代码实体
        if funcType[funcs_id[func_name]]!="dynamic import":
            emb,discode=function_extract(f, callees,callers)
            print("{}\t--{}\n".format(func_name,emb))
            attr[funcs_id[func_name]]=emb

            if sample_extrace == False:
                discodes[funcs_id[func_name]]=discode
                #函数基本块的邻接矩阵、特征向量和对应的汇编代码
                cfg_adj,cfg_attr,cfg_dis=BB_extract(f)                             
                cfg[funcs_id[func_name]]=cfg_adj
                cattr[funcs_id[func_name]]=cfg_attr
                cdis[funcs_id[func_name]]=cfg_dis

    #获得整个函数的邻接矩阵
    cg=cg_extract(funcs_id, callees, func_num)

    #保存图的结点信息
    with open(graphInfLog,"a") as f:
        lth=0
        for temp in cfg:
            lth+=len(temp)
        f.write(str(idaapi.get_root_filename())+"---cfgNum={}  cgNum={}\n".format(lth,len(cg)))

    #保存需要模型训练的信息
    file=shelve.open(dataSavePath+"\\"+str(idaapi.get_root_filename()))
    file["cg"]=cg
    file["cgattr"]=attr
    file["funcType"]=funcType
    file["func_id"]=funcs_id
    if sample_extrace == False:
        file["cfg"]=cfg
        file["cfgattr"]=cattr
        #应该保存谁调用了它，而不是它调用了谁，聚合邻居哇
        file["caller"]=callers
        file["callee"]=callees
        file["discodes"]=discodes
        file["cdiscodes"]=cdis
    file.close()
    print("{}数据提取完成".format(time.time()))

def main():
    q = None   
    f = None
    try:
        idaapi.auto_wait()
        controller()  
        idc.qexit(0)
    except Exception as e:
        with open(extract_error_log, "a+") as file:
            apiName=get_imports_name()
            file.write(apiName+"  "+e+"\n")

if __name__ =="__main__":
    main()

